from prometheus_client import start_http_server
from prometheus_client import Gauge
from utils.commons import *
from utils.statistics import *


def remove_gauge(swap_df, spot_df, _account_info, side: str):
    for item in _account_info[f'{side}_symbol']:
        if item not in swap_df.index and item not in spot_df.index:
            try:
                print('remove ', item)
                _account_info[side].remove(item, 'change')
                _account_info[side].remove(item, 'pos_u')
                _account_info[side].remove(item, 'pnl_u')
                _account_info[side].remove(item, 'avg_price')
                _account_info[side].remove(item, 'cur_price')
            except KeyError:
                pass


def save_gauge(_df, _account_info, _long_arr, _short_arr):
    for symbol, row in _df.iterrows():
        if row['side'] == 1:
            symbol_gauge = _account_info['long']
            _long_arr.append(symbol)
        else:
            symbol_gauge = _account_info['short']
            _short_arr.append(symbol)
        print(f"{symbol} 持仓方向：{row['side']} 仓位：{row['pos_u']}")
        # 写数据
        symbol_gauge.labels(symbol, 'change').set(float(str(row['change']).replace('%', '')))
        symbol_gauge.labels(symbol, 'pos_u').set(row['pos_u'])
        symbol_gauge.labels(symbol, 'pnl_u').set(row['pnl_u'])
        symbol_gauge.labels(symbol, 'avg_price').set(row['avg_price'])
        symbol_gauge.labels(symbol, 'cur_price').set(row['cur_price'])


if __name__ == '__main__':
    start_http_server(9091, addr='0.0.0.0')

    for account_name, account_info in account_config.items():
        account_info['fund'] = Gauge(f'{account_name}_acc', f'{account_name}_acc')
        account_info['long'] = Gauge(f'{account_name}_long', f'{account_name}_long',  ['symbol', 'field'])
        account_info['short'] = Gauge(f'{account_name}_short', f'{account_name}_short', ['symbol', 'field'])

        account_info['long_symbol'] = []
        account_info['short_symbol'] = []

    while True:
        run_time = sleep_until_run_time('1m', if_sleep=True, cheat_seconds=5)  # 设置实时监控账户资金曲线的间隔时间
        for account_name, account_info in account_config.items():
            # ===获取账号的配置
            exchange = account_info['exchange']
            is_use_spot = account_info['is_use_spot']

            try:
                swap_position = get_position_df(exchange)
                swap_equity = get_equity(exchange)
                if is_use_spot:
                    spot_equity, spot_position, spot_usdt, dust_spot = get_spot_position_and_equity(exchange)
                    spot_position = spot_position
                    spot_equity = spot_equity
                else:
                    spot_position = pd.DataFrame()
                    spot_equity = 0
            except:
                traceback.print_exc()
                print(f'当前账号【{account_name}】，获取数据失败')
                continue

            send_spot_df = calc_spot_position(spot_position, account_name)
            send_swap_df = calc_swap_position(swap_position)
            # 清理平仓的labels
            remove_gauge(send_swap_df, send_spot_df, account_info, 'long')
            remove_gauge(send_swap_df, send_spot_df, account_info, 'short')

            # 保存持仓快照Gauge
            long_arr = []
            short_arr = []
            save_gauge(send_spot_df, account_info, long_arr, short_arr)
            save_gauge(send_swap_df, account_info, long_arr, short_arr)

            # 保存这次抓取数据的持仓名单
            account_info['long_symbol'] = long_arr
            account_info['short_symbol'] = short_arr

            print(f'{account_name} 现货净值：{spot_equity}, 合约持仓：{swap_equity}, 合约涨跌幅{swap_position["持仓盈亏"].sum()}')
            current_asset = round(spot_equity + swap_equity + swap_position["持仓盈亏"].sum(), 2)  # 记录账户总净值
            # 保存资金曲线Gauge
            account_info['fund'].set(current_asset)
            # 休息一下，防止超权重
            time.sleep(10)
